package com.mck.controller;

import com.mck.entity.AiRequest;
import com.mck.entity.AiResponse;
import com.mck.service.AiService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.security.concurrent.DelegatingSecurityContextExecutorService;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

import jakarta.annotation.PostConstruct;
import jakarta.annotation.PreDestroy;

@RestController
@RequestMapping("/api/ai")
public class AiController {

  @Autowired
  private AiService aiService;

  private ExecutorService executorService;
  private ExecutorService securityExecutorService;

  @PostConstruct
  public void init() {
    // 创建基础线程池
    this.executorService = Executors.newCachedThreadPool();
    // 创建包装了安全上下文的线程池
    this.securityExecutorService = new DelegatingSecurityContextExecutorService(executorService);
  }

  @PreDestroy
  public void destroy() {
    // 关闭线程池
    if (executorService != null) {
      executorService.shutdown();
    }
  }

  @PostMapping("/chat")
  public AiResponse chat(@RequestBody AiRequest request) {
    return aiService.processAiRequest(request);
  }

  @PostMapping(value = "/chat-stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
  public SseEmitter chatStream(@RequestBody AiRequest request) {
    SseEmitter emitter = new SseEmitter(180000L); // 3分钟超时

    // 添加完成回调，确保资源被正确释放
    emitter.onCompletion(() -> {
      System.out.println("SSE 连接已完成");
    });

    emitter.onTimeout(() -> {
      System.out.println("SSE 连接超时");
    });

    emitter.onError(ex -> {
      System.err.println("SSE 连接发生错误: " + ex.getMessage());
    });

    // 使用包装了安全上下文的执行器服务
    securityExecutorService.execute(() -> {
      try {
        // 发送开始事件
        emitter.send(SseEmitter.event()
            .name("start")
            .data("开始生成回答..."));

        // 调用服务层的流式处理方法
        aiService.processAiRequestStream(request, emitter);

        // 发送完成事件
        emitter.send(SseEmitter.event()
            .name("end")
            .data("回答生成完毕"));

        emitter.complete();
      } catch (Exception e) {
        System.err.println("处理流式请求时出错: " + e.getMessage());
        e.printStackTrace();
        try {
          emitter.send(SseEmitter.event()
              .name("error")
              .data("处理请求时发生错误: " + e.getMessage()));
          emitter.complete();
        } catch (IOException ex) {
          // 忽略发送错误消息时的异常
        }
      }
    });

    return emitter;
  }
}